# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import concurrent.futures
import logging
import math
import sys

from absl import app, flags

from coin.base.datetime_util import to_datetime
from coin.base.param_util import to_list
from coin.strategy.mm.simple_sim.runner_util import (get_profile_module,
                                                     get_num_strategies,
                                                     run_sim_strategy)


class _DebugExecutor:
  def submit(self, fn, *args, **kwargs):
    f = concurrent.futures.Future()
    try:
      res = fn(*args, **kwargs)
      f.set_result(res)
    except Exception as e:
      f.set_exception(e)
    return f

  def __enter__(self):
    return self

  def __exit__(self, exc_type, exc, exc_tb):
    pass


def _get_executor():
  FLAGS = flags.FLAGS
  if FLAGS.debug_executor:
    return _DebugExecutor()
  else:
    return concurrent.futures.ProcessPoolExecutor(max_workers=flags.FLAGS.num_cpu)


def main(argv):
  FLAGS = flags.FLAGS

  try:
    profile_name = argv[1]
    profile_args = argv[2:]
    profile = get_profile_module(profile_name, profile_args)
  except IndexError:
    sys.stderr.write('Usage: python -m %s.runner <profile>\n' % __package__)
    return 1
  except ModuleNotFoundError:
    return 2

  time_ranges = to_list(profile.get_time_ranges())

  futures = []

  with _get_executor() as executor:
    for from_ts, to_ts in time_ranges:
      # TODO(jaewon)
      num_strat = get_num_strategies(profile, from_ts, to_ts)
      num_group = int(math.ceil(num_strat / float(flags.FLAGS.num_strategy_per_group)))
      step = int(math.ceil(num_strat / float(num_group)))

      for i in range(num_group):
        from_index = i * step
        to_index = (min(from_index + step - 1, num_strat - 1) if not FLAGS.test_run else from_index)

        if FLAGS.test_run and len(futures) != 0:
          continue

        logging.info('Initiating %s - %s: %d - %d',
                     to_datetime(from_ts),
                     to_datetime(to_ts),
                     from_index,
                     to_index)
        fut = executor.submit(run_sim_strategy,
                              profile_name,
                              from_ts,
                              to_ts,
                              from_index,
                              to_index,
                              profile_args=profile_args,
                              root_dir=FLAGS.flow_root_dir,
                              from_feed_cache=FLAGS.feed_cache)
        futures.append(fut)

  results = []
  for fut in futures:
    try:
      results += fut.result()
    except Exception:
      logging.exception('Exception is raised')
      continue

  profile.aggregate_result(results)


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  logging.getLogger('coin.exchange.util.feed_checker').setLevel(logging.WARNING)

  flags.DEFINE_bool('feed_cache', False, 'Use feed cache.')

  flags.DEFINE_string('flow_root_dir',
                      None,
                      'If specified, use the specified directory as flow roor directory.')

  flags.DEFINE_bool('test_run', False, 'If true, it run only one strategy profile for testing.')

  flags.DEFINE_integer('num_cpu', 32, 'Number of CPU used to run simulator')

  flags.DEFINE_integer('num_strategy_per_group', 4, 'Number of strategies running in one process.')

  flags.DEFINE_bool('debug_executor', False, 'If true, runs on debug executor')

  app.run(main)
